import numpy as np
import math

class Memory(object):
    def __init__(self, capacity, dims):
        self.capacity = capacity
        self.data = np.zeros((capacity, dims))
        self.pointer = 0

    def store_transition(self, s, a, r, s_):
        transition = np.hstack((s, a, [r], s_))
        index = self.pointer % self.capacity  # replace the old memory with new memory
        self.data[index, :] = transition
        self.pointer += 1

    def sample(self, n):
        assert self.pointer >= self.capacity, '存储尚未填满'
        indices = np.random.choice(self.capacity, size=n)
        return self.data[indices, :]

class preprocess():
    def __init__(self,observation):
        """
        当前为输入敌我双方位置和弹药量，之后要再加入导弹位置
        :param observation:
        """
        self.location_F=observation[0]
        self.location_E=observation[1]
        self.n_ammo=observation[2]
        self.pointing=observation[8]

    def get_input(self):
        dis=self.cal_distance()
        hight=self.cal_diff_hight()
        ang=self.cal_ang()
        pointing=self.cal_pointing()
        ammo=self.n_ammo
        result=np.array([dis,hight,ang,pointing,ammo])
        return result

    def cal_distance(self):
        """
        坐标算距离
        :param a:
        :param b:
        :return:
        """
        a=self.location_F
        b=self.location_E
        diff=(a[0]-b[0])**2+(a[1]-b[1])**2+(a[2]-b[2])**2
        result=math.sqrt(diff)
        return result

    def cal_diff_hight(self):
        a = self.location_F
        b = self.location_E
        return a[2]-b[2]

    def cal_ang(self):
        """
        计算目标点的偏角
        :return:
        """
        a = self.location_F
        b = self.location_E
        y=b[1]-a[1]
        x=b[0]-a[0]
        ang=math.atan2(y,x)
        return ang

    def cal_pointing(self):
        return self.pointing




